#!/usr/bin/env python3
#coding=utf8

import os
import sys
import time
from threading import Thread
from multiprocessing import Pool,Queue
from multiprocessing.managers import SyncManager
from multiprocessing.managers import BaseManager
from subprocess import Popen,PIPE
from core.models import Task,Plugin

sys.path.append('plugins')

def call(args):
    task = args[0]
    plugin_name = args[1]
    plugin_name = plugin_name.strip()
    if plugin_name.endswith(('.py','.PY')):
        plugin_name = plugin_name[:-3]
    task.task_pid = os.getpid()
    task.task_state = 1 #进行中
    task.starttime = int(time.time())
    task.save()
    try:
        ret = __import__(plugin_name).main(task.task_id)
        task.task_state = 2 #已完成
    except Exception as e:
        print(e)
        task.task_state = 3 #出错
    task.endtime = int(time.time())
    task.save()

class TaskServer(Thread):
    def __init__(self):
        Thread.__init__(self)
        self.pool = Pool()

    def run(self):
        while True:
            print('[+] Wait Task %s'%time.time(),end='\r')
            tasklist = []
            for task in (Task.select().where(Task.task_state==0)):
                tasklist.append([task,task.task_type.plugin_path])
            self.pool.map(call,tasklist)
            time.sleep(3)

class RpcServer(Thread):
    class RpcManageServer(SyncManager):pass
    def __init__(self,address=('127.0.0.1',8315),key='KEY'):
        Thread.__init__(self)
        self.queue = Queue()
        self.RpcManageServer.register('get_queue', callable=lambda:self.queue)
        self.manager = self.RpcManageServer(address, authkey=key.encode()).get_server()

    def run(self):
        print("RpcServer Runing ...")
        self.manager.serve_forever()

class RpcHandler(Thread):
    class RpcManageHandler(SyncManager):pass
    def __init__(self,address=('127.0.0.1',8315),key='KEY'):
        Thread.__init__(self)
        self.RpcManageHandler.register('get_queue')
        self.manager = self.RpcManageHandler(address, authkey=key.encode())
        self.manager.connect()

    def run(self):
        result_queue = self.manager.get_queue()
        while True:
            res = result_queue.get()
            print(res)

class RpcClient(Thread):
    def __init__(self,address=('127.0.0.1', 8315),key='KEY'):
        Thread.__init__(self)
        BaseManager.register('get_queue')
        conn = BaseManager(address, authkey=key.encode())
        conn.connect()
        self.queue = conn.get_queue()

    def save(self,data):
        self.queue.put(data)


if __name__ == '__main__':
    ts = TaskServer()
    ts.start()
    #rs = RpcServer()
    #rs.start()
    #rh = RpcHandler()
    #rh.start()
    #rc = RpcClient()
    #rc.start()
